import matplotlib.pyplot as plt
import numpy as np
import os
import argparse
import shutil


def plotAirfoilIter(coordinates_list, file_path, file_list):
    '''
        在同一张图上展示一系列翼型
    :param coordinates:
    :param file_path:
    :return:
    '''

    plt.rc('font', family='Latin Modern Roman')

    crd_num = len(coordinates_list)
    linestyle_list = ['solid',
                      ':', '--', '-.',
                      ':', '--', '-.']
    color_list = ["black", "red", "blue", "red", "blue", "red", "blue"]
    long = 7
    height = 6.6

    fig = plt.figure(figsize=(long, height), dpi=80)
    fig.canvas.draw()

    for i, coordinates in enumerate(coordinates_list):
        plt.plot(coordinates[:, 0], coordinates[:, 1],
                 linewidth=1,
                 linestyle=linestyle_list[i],
                 label=file_list[i][9:-4],
                 c=color_list[i])

    font = {'weight': 'normal',
             'size': 16,
             }

    # 横纵坐标标题
    plt.xlabel("Iteration steps", font)
    plt.ylabel(r"$log_{10}f$", font)
    # plt.ylabel("Objective function value of cp")

    font1 = 13
    # 横纵坐标刻度
    plt.xticks(size=font1)
    plt.yticks(size=font1)

    # plt.xlim((-0.01, 1.01))
    # plt.ylim((-1.15, 1.6))
    plt.legend(fontsize=font1)
    # plt.grid()
    plt.yscale("log")

    plt.savefig(file_path, dpi=1000, bbox_inches='tight')
    plt.close()


if __name__ == "__main__":

    '''
        -r /home/linjx/pycharmworkspace/visual/visual_iteration/airfoil_cp_mlcls -s /home/linjx/pycharmworkspace/visual/visual_airfoil/plot
        -r /home/linjx/pycharmworkspace/visual/visual_iteration/airfoil_cp_clsVSnotcls -s /home/linjx/pycharmworkspace/visual/visual_airfoil/plot
    '''

    # 命令行参数
    # parser = argparse.ArgumentParser()
    # parser.add_argument("-r", "--read_dir", type=str, help="待处理翼型所在文件夹路径", required=True)
    # parser.add_argument("-s", "--save_dir", type=str, help="处理后的翼型保存路径", required=True)
    # args = parser.parse_args()


    # 反设计
    InvOrPst = r"InverseDesign"
    pic_name_list = ["iterations_mlcls", "iterations_clsVSnotcls"]

    # 正设计
    # pic_name_list = ["iterations_cd_clsVSnotcls"]
    # InvOrPst = r"PositiveDesign"


    read_dir_list = ["/home/linjx/pycharmworkspace/visual/visual_iteration/{}/{}".format(InvOrPst, pic_name)
                     for pic_name in pic_name_list]
    save_dir = "/home/linjx/pycharmworkspace/visual/visual_iteration/plot"
    save_pic_list = ["{}.png".format(pic_name) for pic_name in pic_name_list]

    # if os.path.exists(save_dir):
    #     shutil.rmtree(save_dir)
    # os.mkdir(save_dir)

    for read_dir, save_pig in zip(read_dir_list, save_pic_list):
        # 画多个翼型
        save_path = os.path.join(save_dir, save_pig)
        coordinates_list = []
        file_list = []

        # 文件名排序
        filenames = os.listdir(read_dir)
        filenames.sort(key=lambda x: int(x[0:2]))

        for file in filenames:
            print(file)
            file_path = os.path.join(read_dir, file)
            coordinates = np.loadtxt(file_path, skiprows=1, usecols=(0, 1))[:-3]
            coordinates_list.append(coordinates)
            file_list.append(file)
        plotAirfoilIter(coordinates_list, save_path, file_list)
        print("翼型优化迭代图已保存在{}".format(save_path))

